package org.rayson.share.transport;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;

import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;
import javax.net.ssl.SSLEngineResult.HandshakeStatus;
import javax.net.ssl.SSLException;

public class SslChannel extends NioChannel {

	private final NioSelector selector;
	private final SSLEngine engine;
	private final ByteBuffer appIn, netIn, netOut;
	private boolean handshaked;
	private static final int EOF = -1;

	public SslChannel(final SelectionKey sk, final SocketChannel sc,
			final NioSelector selector, final SSLEngine engine) {
		super(sk, sc);
		this.selector = selector;
		this.engine = engine;
		appIn = ByteBuffer.allocate(engine.getSession()
				.getApplicationBufferSize());
		netIn = ByteBuffer.allocate(engine.getSession().getPacketBufferSize());
		netOut = ByteBuffer.allocate(engine.getSession().getPacketBufferSize());
	}

	@Override
	public int read(final ByteBuffer dst) throws IOException {

		if (sc.read(netIn) == EOF)
			return EOF;

		netIn.flip();
		final SSLEngineResult result = engine.unwrap(netIn, appIn);
		netIn.compact();

		if (!handshaked) {
			OUTER: while (true) {
				final HandshakeStatus hs = engine.getHandshakeStatus();
				switch (hs) {
				case FINISHED:
				case NOT_HANDSHAKING:
					handshaked = true;
					break OUTER;
				case NEED_TASK:

					engine.getDelegatedTask().run();
					continue;

				case NEED_UNWRAP:

					this.selector.addReadEvent(sk);
					return 0;

				case NEED_WRAP:
					this.selector.addWriteEvent(sk);
					return 0;

				default:
					throw new SSLException("Unknown handshake status");
				}
			}
		}

		final int maxRead = Math.min(dst.remaining(), appIn.position());
		if (maxRead == 0)
			return 0;

		appIn.flip();
		for (int i = 0; i < maxRead; i++) {
			dst.put(appIn.get());
		}
		appIn.compact();

		if (netIn.position() > 0 || appIn.position() > 0) {
			this.selector.addReadEvent(sk);
		}

		return maxRead;
	}

	@Override
	public int write(final ByteBuffer src) throws IOException {

		final SSLEngineResult result = engine.wrap(src, netOut);

		if (!handshaked) {
			OUTER: while (true) {
				final HandshakeStatus hs = engine.getHandshakeStatus();
				switch (hs) {
				case FINISHED:
				case NOT_HANDSHAKING:
					handshaked = true;
					break OUTER;
				case NEED_TASK:

					engine.getDelegatedTask().run();
					continue;

				case NEED_UNWRAP:

					this.selector.addReadEvent(sk);
					break OUTER;

				case NEED_WRAP:
					this.selector.addWriteEvent(sk);
					break OUTER;

				default:
					throw new SSLException("Unknown handshake status");
				}
			}
		}

		netOut.flip();
		sc.write(netOut);
		netOut.compact();

		return result.bytesConsumed();
	}

}